package nachos.vm;


import java.util.Random;

import nachos.machine.*;
import nachos.threads.KThread;
import nachos.userprog.*;

/**
 * A <tt>UserProcess</tt> that supports demand-paging.
 */
public class VMProcess extends UserProcess {
    /**
     * Allocate a new process.
     */
    public VMProcess() {
    	super();
    	tlbManager = new TLBManager(this.processID);
    }
    
    /**
     * Handle the exec(char *file, int argc, char *argv[]) system call
     * 
     * @return	child process id, or -1 on error 
     */
    private int handleExec(int fileNameAddr, int argCount, int argAddr){
    	System.out.println("handled by VM!");
    	
    	// Get the filename
    	String filename = this.readVirtualMemoryString(fileNameAddr, maxLen);
    	if (filename == null){
    		System.out.println("exec wrong :  + filename = null ");
    		return -1;
    	}
    	
    	// Check file exists
    	OpenFile ftry = VMKernel.fileSystem.open(filename, false);
    	if (ftry == null){
    		System.out.println("exec wrong : " + filename + " open fail!");
    		return -1;
    	}
    	ftry.close();
    	
    	// Create a new Process
    	VMProcess up = new VMProcess();
    	up.parentProcess = this;
    	    	
    	// Get the arguments
    	String []arguments = new String[argCount];
    	byte[] pointer = new byte[4];
    	for (int i = 0; i < argCount; i ++){
    		// read in a pointer
    		this.readVirtualMemory(argAddr, pointer);
    		int intPointer = Lib.bytesToInt(pointer, 0);
    		Lib.debug(dbgProcess, "argAddr = " + argAddr + "   intPointer = " + intPointer);
    		argAddr += 4;
    		// read real content
    		String argument = this.readVirtualMemoryString(intPointer, maxLen);
    		if (argument != null){
    			arguments[i] = argument;
    			Lib.debug(dbgProcess, "\targ[" + i + "] = " + arguments[i] + 
        				" length = " + arguments[i].length());
    		}
    		
    	}
    	
    	// Execute it !
    	if (up.execute(filename, arguments)){
    		// Return the ID of the new process
    		Lib.debug(dbgProcess, "process id = " + up.processID);
    		return up.processID;
    	}
    	else{
    		//up.handleExit(0);
    		return -1;
    	}
    }

    /**
     * Execute the specified program with the specified arguments. Attempts to
     * load the program, and then forks a thread to run it.
     *
     * @param	name	the name of the file containing the executable.
     * @param	args	the arguments to pass to the executable.
     * @return	<tt>true</tt> if the program was successfully executed.
     */
    public boolean execute(String name, String[] args) {
		if (!load(name, args)){
			System.out.println("load fail!");
		    return false;
		}
		
		thread = new UThread(this);
		thread.setName(name).fork();
		
		tlbManager.setAllItemInvalid();
		
		return true;
    }
    
    
    /**
     * Save the state of this process in preparation for a context switch.
     * Called by <tt>UThread.saveState()</tt>.
     */
    public void saveState() {
    	super.saveState();
    	tlbManager.setAllItemInvalid();
    	Lib.debug(dbgVM, " save state, pid = " + this.processID + " all TLB items invalid");
    }

    /**
     * Restore the state of this process after a context switch. Called by
     * <tt>UThread.restoreState()</tt>.
     */
    public void restoreState() {
	}
    
    /**
     * Transfer data from this process's virtual memory to the specified array.
     * This method handles address translation details. This method must
     * <i>not</i> destroy the current process if an error occurs, but instead
     * should return the number of bytes successfully copied (or zero if no
     * data could be copied).
     *
     * @param	vaddr	the first byte of virtual memory to read.
     * @param	data	the array where the data will be stored.
     * @param	offset	the first byte to write in the array.
     * @param	length	the number of bytes to transfer from virtual memory to
     *			the array.
     * @return	the number of bytes successfully transferred.
     */
    public int readVirtualMemory(int vaddr, byte[] data, int offset, int length) {
		Lib.assertTrue(offset >= 0 && length >= 0 && offset+length <= data.length);
		
		Lib.debug(dbgProcess, "pid = " + this.processID + " read virtual mem : vaddr = " + vaddr + " length = " + length);		
	
		byte[] memory = Machine.processor().getMemory();

		int amount = 0, tamount;// total size read
		int vpn, ppn, srcOffset, desOffset = offset;
		
		// Read the first page
		vpn = vaddr/pageSize;
		ppn = -1;
		for (int i = 0; i < Machine.processor().getTLBSize(); i ++){
			if (Machine.processor().readTLBEntry(i).valid 
					&& Machine.processor().readTLBEntry(i).vpn == vpn){
				ppn = Machine.processor().readTLBEntry(i).ppn;
			}
		}
		// if cannot find this page in TLB
		if (ppn == -1){
			ppn = tlbManager.request(vpn);
		}
		else{
			// inform this page been used
			PagingManager.informUsed(ppn);
		}
		
		PagingManager.lockPage(ppn);
		
		srcOffset = vaddr%pageSize + ppn*pageSize;
		tamount = Math.min(pageSize - vaddr%pageSize, length);
		Lib.debug(dbgProcess, "read virtual memory, memory length = " + memory.length + 
				"   srcOffset = " + srcOffset + 
				"   desOffset = " + desOffset + 
				"   tamount = " + tamount);
		System.arraycopy(memory, srcOffset, data, desOffset, tamount);
		PagingManager.unlockPage(ppn);
		
		amount += tamount;
		desOffset += tamount;
		informPageUsed(vpn);
		tlbManager.setUsed(vpn);
		
		vpn++;
		
		
		// Read pages in the medium
		while(vpn < (vaddr+length)/pageSize){			
			ppn = -1;
			for (int i = 0; i < Machine.processor().getTLBSize(); i ++){
				if (Machine.processor().readTLBEntry(i).valid 
						&& Machine.processor().readTLBEntry(i).vpn == vpn){
					ppn = Machine.processor().readTLBEntry(i).ppn;
				}
			}
			// if cannot find this page in TLB
			if (ppn == -1){
				ppn = tlbManager.request(vpn);
			}
			else{
				// inform this page been used
				PagingManager.informUsed(ppn);
			}
			
			PagingManager.lockPage(ppn);
			
			srcOffset = ppn*pageSize;
			tamount = pageSize;
			Lib.debug(dbgProcess, "srcOffset = " + srcOffset + "   desOffset = " + desOffset + "   tamount = " + tamount);
			System.arraycopy(memory, srcOffset, data, desOffset, tamount);
			
			PagingManager.unlockPage(ppn);
			
			amount += tamount;
			desOffset += tamount;
			informPageUsed(vpn);
			tlbManager.setUsed(vpn);
			vpn++;
		}
		
		// Read the last page
		if (vpn == (vaddr+length)/pageSize){
			informPageUsed(vpn);
			ppn = -1;
			for (int i = 0; i < Machine.processor().getTLBSize(); i ++){
				if (Machine.processor().readTLBEntry(i).valid 
						&& Machine.processor().readTLBEntry(i).vpn == vpn){
					ppn = Machine.processor().readTLBEntry(i).ppn;
				}
			}
			// if cannot find this page in TLB
			if (ppn == -1){
				ppn = tlbManager.request(vpn);
			}
			else{
				// inform this page been used
				PagingManager.informUsed(ppn);
			}
			
			PagingManager.lockPage(ppn);
			
			srcOffset = ppn * pageSize;
			tamount = (vaddr+length)%pageSize;
			Lib.debug(dbgProcess, "srcOffset = " + srcOffset + "   desOffset = " + desOffset + "   tamount = " + tamount);
			System.arraycopy(memory, srcOffset, data, desOffset, tamount);
			
			PagingManager.unlockPage(ppn);
			
			amount += tamount;
			tlbManager.setUsed(vpn);
		}
		
		return amount;
		
    }
    
    /**
     * Transfer data from the specified array to this process's virtual memory.
     * This method handles address translation details. This method must
     * <i>not</i> destroy the current process if an error occurs, but instead
     * should return the number of bytes successfully copied (or zero if no
     * data could be copied).
     *
     * @param	vaddr	the first byte of virtual memory to write.
     * @param	data	the array containing the data to transfer.
     * @param	offset	the first byte to transfer from the array.
     * @param	length	the number of bytes to transfer from the array to
     *			virtual memory.
     * @return	the number of bytes successfully transferred.
     */
    public int writeVirtualMemory(int vaddr, byte[] data, int offset, int length) {
    	Lib.assertTrue(offset >= 0 && length >= 0 && offset+length <= data.length);
	
    	Lib.debug(dbgProcess, "write virtual mem : vaddr = " + vaddr + " length = " + length);
    	
		byte[] memory = Machine.processor().getMemory();
		
		int amount = 0, tamount;// total size write		
		int vpn, ppn, srcOffset = offset, desOffset;
		
		// write the first page
		vpn = vaddr/pageSize;
		ppn = -1;
		for (int i = 0; i < Machine.processor().getTLBSize(); i ++){
			if (Machine.processor().readTLBEntry(i).valid 
					&& Machine.processor().readTLBEntry(i).vpn == vpn){
				ppn = Machine.processor().readTLBEntry(i).ppn;
			}
		}
		// if cannot find this page in TLB
				if (ppn == -1){
					ppn = tlbManager.request(vpn);
				}
				else{
					// inform this page been used
					PagingManager.informUsed(ppn);
				}
		
		while (PagingManager.isLocked(ppn)){
			System.out.println("write locked page");
			KThread.yield();
		}
			
		PagingManager.lockPage(ppn);
		
		desOffset = vaddr%pageSize + ppn*pageSize;
		tamount = Math.min(pageSize - vaddr%pageSize, length);
		Lib.debug(dbgProcess, "write virtual memory, srcOffset = " + srcOffset + "   desOffset = " + desOffset + "   tamount = " + tamount);
		System.arraycopy(data, srcOffset, memory, desOffset, tamount);		
		this.informPageDirty(ppn);
		
		//Machine.interrupt().restore(intStat);
		PagingManager.unlockPage(ppn);
		
		amount += tamount;
		srcOffset += tamount;
		informPageDirty(vpn);
		tlbManager.setDirty(vpn);
		tlbManager.setUsed(vpn);
		vpn++;
		
		// Write pages in the medium
		while(vpn < (vaddr+length)/pageSize){
			ppn = -1;
			for (int i = 0; i < Machine.processor().getTLBSize(); i ++){
				if (Machine.processor().readTLBEntry(i).valid 
						&& Machine.processor().readTLBEntry(i).vpn == vpn){
					ppn = Machine.processor().readTLBEntry(i).ppn;
				}
			}
			// if cannot find this page in TLB
			if (ppn == -1){
				ppn = tlbManager.request(vpn);
			}
			else{
				// inform this page been used
				PagingManager.informUsed(ppn);
			}
			
			while (PagingManager.isLocked(ppn)){
				System.out.println("write locked page");
				KThread.yield();
			}
			PagingManager.lockPage(ppn);
			
			desOffset = ppn*pageSize;
			tamount = pageSize;
			Lib.debug(dbgProcess, "srcOffset = " + srcOffset + "   desOffset = " + desOffset + "   tamount = " + tamount);
			System.arraycopy(data, srcOffset, memory, desOffset, tamount);
			this.informPageDirty(ppn);
			
			PagingManager.unlockPage(ppn);
			
			amount += tamount;
			srcOffset += tamount;
			informPageDirty(vpn);
			tlbManager.setDirty(vpn);
			tlbManager.setUsed(vpn);
			vpn++;
		}
		
		// Write the last page
		if (vpn == (vaddr+length)/pageSize){
			informPageDirty(vpn);
			ppn = -1;
			for (int i = 0; i < Machine.processor().getTLBSize(); i ++){
				if (Machine.processor().readTLBEntry(i).valid 
						&& Machine.processor().readTLBEntry(i).vpn == vpn){
					ppn = Machine.processor().readTLBEntry(i).ppn;
				}
			}
			// if cannot find this page in TLB
			if (ppn == -1){
				ppn = tlbManager.request(vpn);
			}
			else{
				// inform this page been used
				PagingManager.informUsed(ppn);
			}
			
			while (PagingManager.isLocked(ppn)){
				System.out.println("write locked page");
				KThread.yield();
			}
			PagingManager.lockPage(ppn);
			
			desOffset = ppn * pageSize;
			tamount = (vaddr+length)%pageSize;
			Lib.debug(dbgProcess, "srcOffset = " + srcOffset + "   desOffset = " + desOffset + "   tamount = " + tamount);
			System.arraycopy(data, srcOffset, memory, desOffset, tamount);
			this.informPageDirty(ppn);
			
			PagingManager.unlockPage(ppn);
			
			amount += tamount;
			tlbManager.setDirty(vpn);
			tlbManager.setUsed(vpn);
		}
		return amount;
    }

    /**
     * Initializes page tables for this process so that the executable can be
     * demand-paged.
     *
     * @return	<tt>true</tt> if successful.
     */
    protected boolean loadSections() {
		// just initialize sections info, load later
		for (int s = 0; s < coff.getNumSections(); s++) {
		    CoffSection section = coff.getSection(s);		    
		    for (int i = 0; i < section.getLength(); i++) {
				int vpn = section.getFirstVPN()+i;
				PagingManager.lazy(this.processID, vpn, section, i, true);
		    }
		}
		
		// Load the stack pages
		int vpn = initialSP / pageSize -1;
		for (int i = 0; i < stackPages; i ++){
			PagingManager.lazy(processID, vpn, null, -1, false);			
			vpn--;
		}
		
		// Load the last page for the arguments
		vpn = numPages - 1;
		PagingManager.lazy(processID, vpn, null, -1, false);
		
		return true;
    }

    /**
     * Release any resources allocated by <tt>loadSections()</tt>.
     */
	protected void unloadSections() {
		Lib.debug(VMProcess.dbgVM, "unload sections");
		SwapManager.releaseProcess(processID);
		coff.close();
	}    

    /**
     * Handle a user exception. Called by
     * <tt>UserKernel.exceptionHandler()</tt>. The
     * <i>cause</i> argument identifies which exception occurred; see the
     * <tt>Processor.exceptionZZZ</tt> constants.
     *
     * @param	cause	the user exception that occurred.
     */
    public void handleException(int cause) {	
		switch (cause) {		
		case Processor.exceptionTLBMiss:
			handleTLBMiss();
			break;
		default:
		    super.handleException(cause);
		    break;
		}
    }
    
    private static final int syscallExec = 2;
    public int handleSyscall(int syscall, int a0, int a1, int a2, int a3) {
    	switch (syscall) {
    	case syscallExec:
			return handleExec(a0, a1, a2);
		default:
				return super.handleSyscall(syscall, a0, a1, a2, a3);
    	}
    }
   
    /**
     * Handle TLB miss 
     */
    private void handleTLBMiss() {
    	//tlbManager.printTLBStatus(true, 0);
    	
		// First, get the virtual address causing the TLB miss
    	// and calculate the virtual page number
    	int vAddr = Machine.processor().readRegister(Processor.regBadVAddr);
    	int vpn = vAddr / pageSize;
    
    	// Then, let TLBManager handle this problem :)
    	tlbManager.request(vpn);
	}
    
    @Override
    protected void informPageUsed(int vpn){
    	tlbManager.setUsed(vpn);
    }
    
    @Override
    protected void informPageDirty(int vpn){
    	tlbManager.setDirty(vpn);
    }
    
    class TLBManager{
    	TLBManager(int _pid){
    		pid = _pid;
    	}
    	
    	public void printTLBStatus(boolean printAll, int index){
    		if (printAll){
    			for (int i = 0; i < Machine.processor().getTLBSize(); i ++){
    				TranslationEntry te = Machine.processor().readTLBEntry(i); 
    				Lib.debug(VMProcess.dbgVM, i + " : vpn = " + te.vpn + 
    						" ppn = " + te.ppn + " used = " + te.used + 
    						" dirty = " + te.dirty + " valid = " + te.valid + 
    						" read-only = " + te.readOnly);
    			}
    		}
    		else{
    			TranslationEntry te = Machine.processor().readTLBEntry(index);
    			Lib.debug(VMProcess.dbgVM, index + " : vpn = " + te.vpn + 
    					" ppn = " + te.ppn + " used = " + te.used + 
    					" dirty = " + te.dirty + " valid = " + te.valid + 
    					" read-only = " + te.readOnly);
    		}
    		return;
    	}
    	
    	// handle TLBMiss 
    	int request(int vpn){
    		// look up the page table
    		Lib.debug(dbgVM, "TLB miss! pid = " + pid + " vpn = " + vpn);
    		//this.printTLBStatus(true, 0);
    		TranslationEntry te = PagingManager.lookUp(pid, vpn);
    		
    		Lib.debug(dbgVM, "tlb got entry from page table : vpn = " + te.vpn + " ppn = " + te.ppn);
    		
    		// update the TLB
    		writeEntry(te);
    		
    		return te.ppn;
    	}
    	
    	// write a new Entry into the TLB
    	void writeEntry(TranslationEntry te){
    		// calculate the ppn, make a new Entry. 
    		TranslationEntry e = new TranslationEntry();
    		e.ppn = te.ppn;
    		e.vpn = te.vpn;
    		e.readOnly = te.readOnly;
    		e.used = true;
    		e.valid = true;
    		e.dirty = false;
    		
    		// find a position, may cause another Entry invalid.
    		//tlbManager.printTLBStatus(true, 0);
    		int pos = -1;
    		for (int i = 0; i < Machine.processor().getTLBSize(); i ++){
    			if (Machine.processor().readTLBEntry(i).valid == false){
    				pos = i;
    				break;
    			}
    		}
    		if (pos == -1){
    			pos = new Random().nextInt(Machine.processor().getTLBSize());
    		}
    		
    		// if this position is dirty, inform pageTable
    		if (Machine.processor().readTLBEntry(pos).dirty){
    			PagingManager.informDirty(pos);    			
    		}
    		
    		// write in.
			Machine.processor().writeTLBEntry(pos, e);
			//Lib.debug(dbgVM, "write new TLB : pos = " + pos + " ppn = " + e.ppn + " vpn = " + e.vpn);
			//tlbManager.printTLBStatus(true, 0);
    	}
    	
    	void setAllItemInvalid(){
    		for (int i = 0; i < Machine.processor().getTLBSize(); i ++){
    			TranslationEntry te = Machine.processor().readTLBEntry(i);
    			te.valid = false;
    			Machine.processor().writeTLBEntry(i, te);
    		}
    	}
    	
    	/**
    	 * Set this page used. Might be called by readVirtualMemory.
    	 * @param vpn
    	 */
    	void setUsed(int vpn){
    		for (int i = 0; i < Machine.processor().getTLBSize(); i ++){
    			TranslationEntry te = Machine.processor().readTLBEntry(i);
    			
    			// if vpn address match
    			if (te.vpn == vpn){
    				// modify the used bit, and write back
    				te.used = true;
    				Machine.processor().writeTLBEntry(i, te);
    			}
    		}
    	}
    	
    	/**
    	 * Set this page dirty. Might be called by writeVirtualMemory.
    	 * @param vpn
    	 */
    	void setDirty(int vpn){
    		for (int i = 0; i < Machine.processor().getTLBSize(); i ++){
    			TranslationEntry te = Machine.processor().readTLBEntry(i);
    			
    			// if vpn address match
    			if (te.vpn == vpn){
    				// modify the dirty bit, and write back
    				te.dirty = true;
    			}
    		}
    	}
    	
    	/** process id*/
    	int pid;
    	
    }
    
    private TLBManager tlbManager;

	private static final int pageSize = Processor.pageSize;
    private static final char dbgProcess = 'a';
    public static final char dbgVM = 'v';
}
